import cache

class Corpus(object):
    def __init__(self, file, language, add=False):
        self.language = language
        self.corpus = []
        self.unigram = cache.Cache('unigram_%s' % language, dict())
        self.bigram = cache.Cache('bigram_%s' % language, dict())
        self.trigram = cache.Cache('trigram_%s' % language, dict())
        self.allignment_count = cache.Cache('allignment_count_%s' % language, dict())
        self.load_corpus(file)
        if add:
            self.add_null_word()
        self.load_statistics()

    def add_null_word(self):
        """ adds dummy NULLL word to the start of each sentence """
        for i in xrange(len(self.corpus)):
            self.corpus[i] = ['NULL'] + self.corpus[i]

    def load_corpus(self, file):
        f = open(file)
        sentences = f.readlines()
        for sentence in sentences:
            sentence = sentence.replace('\n', '').split(' ')
            self.corpus.append(sentence)        
        f.close()

    def load_statistics(self):
        if len(self.unigram.cache) == 0:
            self.fill_ngram(1, self.unigram.cache)
            self.unigram.save()
        if len(self.bigram.cache) == 0:
            self.fill_ngram(2, self.bigram.cache)
            self.bigram.save()
        if len(self.trigram.cache) == 0:    
            self.fill_ngram(3, self.trigram.cache)
            self.trigram.save()

    def fill_ngram(self, n, ngram_dict):
        for sentence in self.corpus:
            sentence = Corpus.add_sentence_delimiters(sentence)
            for i in range(0, len(sentence)-(n-1)):
                key = ' '.join(sentence[i:i+n])
                if key in ngram_dict:
                    ngram_dict[key] = ngram_dict.get(key) + 1
                else:
                    ngram_dict[key] = 1
            
            
    def size(self):
        return len(self.corpus)       

    def get_total_word_count(self):
        return sum(self.unigram.cache.values())

    def get_allignment_count(self, corpus_f, e, f=None):
        key = (e, f) if f else e
        if key in self.allignment_count.cache:
            return self.allignment_count.cache.get(key)
        else:
            if f and (f, e) in corpus_f.allignment_count.cache:
                count = corpus_f.allignment_count.cache.get((f, e))
            else:
                count = self.count_allignment(corpus_f, e, f)
            self.allignment_count.cache[key] = count
        return count

    def count_allignment(self, corpus_f, e, f=None):
        count = 0
        for i, s in enumerate(self.corpus):
            positions = [p for p, w in enumerate(s) if w == e]
            for p in positions:
                if ((len(corpus_f.corpus[i]) > p) and 
                        (not f or (corpus_f.corpus[i][p] == f))):
                    count += 1
        return count


    @staticmethod
    def add_sentence_delimiters(sentence):
        return ['<s>'] + sentence + ['<\s>']
